package org.komarichin.filters.trim;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.servlet.ServletOutputStream;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStreamWriter;
import java.io.PrintWriter;

/**
 * User: Mykhaylo Komarichyn
 * Date: 18.02.2010
 * Time: 18:15:34
 */
public class TrimResponseWrapper extends HttpServletResponseWrapper {

    private static final Logger log = LoggerFactory.getLogger(TrimResponseWrapper.class);

    private ByteArrayOutputStream output;
    private PrintWriter writer;
    private final HttpServletResponse origResponse;
    private int contentLength;
    private String contentType;

    public TrimResponseWrapper(HttpServletResponse httpservletresponse) {
        super(httpservletresponse);
        output = null;
        writer = null;
        output = new ByteArrayOutputStream();
        origResponse = httpservletresponse;
    }

    public byte[] getData() {
        try {
            finishResponse();
        }
        catch (Exception e) {
            log.error(e.getMessage(), e);
        }
        return output.toByteArray();
    }

    public ServletOutputStream getOutputStream() {
        return new FilterServletOutputStream(output);
    }

    public ServletOutputStream createOutputStream() {
        return new FilterServletOutputStream(output);
    }

    public PrintWriter getWriter()
            throws IOException {
        if (writer != null) {
            return writer;
        } else {
            ServletOutputStream servletoutputstream = getOutputStream();
            writer = new PrintWriter(new OutputStreamWriter(servletoutputstream, origResponse.getCharacterEncoding()));
            return writer;
        }
    }

    public void finishResponse() {
        try {
            if (writer != null) {
                writer.close();
            } else if (output != null) {
                output.flush();
                output.close();
            }
        }
        catch (IOException e) {
            log.error(e.getMessage(), e);
        }
    }

    public void flushBuffer()
            throws IOException {
        output.flush();
    }

    public void setContentLength(int i) {
        contentLength = i;
        super.setContentLength(i);
    }

    public int getContentLength() {
        return contentLength;
    }

    public void setContentType(String s) {
        contentType = s;
        super.setContentType(s);
    }

    public String getContentType() {
        return contentType;
    }
}
